package org.rbio.paper.methylglm.annotate;

import java.io.BufferedReader;
import java.io.BufferedWriter;
import java.io.FileReader;
import java.io.FileWriter;
import java.util.ArrayList;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.TreeMap;

// annotate the DMR
public class DMRAnnotate {
	// annotate DMRs
	public static void annotate(String cpg_ann_file, String dmr_file,
			String out_file) throws Exception {
		// make the CpG annotation table
		Map<String, CpGAnn> ann_tab = new TreeMap<String, CpGAnn>();
		BufferedReader ibuf = new BufferedReader(new FileReader(cpg_ann_file));
		String line = ibuf.readLine();
		while ((line = ibuf.readLine()) != null) {
			String arr[] = line.split("\t");
			ann_tab.put(arr[0] + ":" + arr[1], new CpGAnn(line));
		}
		ibuf.close();

		// annotate the dmrs
		BufferedWriter obuf = new BufferedWriter(new FileWriter(out_file));
		obuf.write("ref\tstart\tend\tcnt\tpos\ttss\tgene\texon\tisl\n");
		ibuf = new BufferedReader(new FileReader(dmr_file));
		line = ibuf.readLine();
		while ((line = ibuf.readLine()) != null) {
			obuf.write(line + "\t" + annotateDMR(line, ann_tab) + "\n");
		}
		ibuf.close();
		obuf.close();
	}

	// annotate the dmr
	public static String annotateDMR(String dmr_str, Map<String, CpGAnn> ann_tab) {
		// parse the dmr line
		String arr[] = dmr_str.split("\t");

		// fill the CpGAnn list
		List<CpGAnn> cpg_ann_lst = new ArrayList<CpGAnn>();
		for (String pos : arr[4].split(",")) {
			cpg_ann_lst.add(ann_tab.get(arr[0] + ":" + pos));
		}

		// make the union of CpG annotations
		boolean isl = false;
		Map<String, GeneDMRAnn> gene_map = new LinkedHashMap<String, GeneDMRAnn>();
		for (CpGAnn ann : cpg_ann_lst) {
			// check if this CpG overlaps an island
			if (ann.isl == true) {
				isl = true;
			}

			// union of CpG genes for this DMR
			for (GeneAnn gene : ann.gene_lst) {
				GeneDMRAnn gene_ann = gene_map.get(gene.name);
				if (gene_ann == null) {
					gene_ann = new GeneDMRAnn(gene);
					gene_map.put(gene.name, gene_ann);
				}
				gene_ann.update(gene);
			}
		}

		// make the result string
		String tss = "NA", gname = "NA", exon = "NA", isl_str = "NA";
		if (isl == true) {
			isl_str = "YES";
		}
		for (GeneDMRAnn gene : gene_map.values()) {
			// tss
			if (gene.in_tss == true) {
				if (tss.compareTo("NA") == 0) {
					tss = gene.name;
				} else {
					System.out.println(dmr_str);
					tss += "," + gene.name;
				}
			}

			// gene
			if (gene.exon.compareTo("NA") != 0) {
				if (gname.compareTo("NA") == 0) {
					gname = gene.name;
					exon = gene.exon;
				} else {
					gname += "," + gene.name;
					exon += "," + gene.exon;
				}
			}
		}
		return tss + "\t" + gname + "\t" + exon + "\t" + isl_str;
	}

	// annotation for a single CpG
	public static class CpGAnn {
		public List<GeneAnn> gene_lst = new ArrayList<GeneAnn>();
		public boolean isl = false;

		// constructor
		public CpGAnn(String line) {
			// parse CpG line
			String arr[] = line.split("\t");
			int len = arr.length;
			isl = (arr[len - 1].compareTo("YES") == 0) ? true : false;
			String exon_arr[] = arr[len - 2].split(",");
			String gene_arr[] = arr[len - 3].split(",");
			String tss_arr[] = arr[len - 4].split(",");

			// make the gene
			if (arr[len - 3].compareTo("NA") != 0) {
				// CpG in at least one gene body
				for (int i = 0; i < gene_arr.length; i++) {
					// fill gene list with exon status
					gene_lst.add(new GeneAnn(gene_arr[i], exon_arr[i]));
				}
			}

			// fill the tss flag for the existing genes, add new genes if the
			// cpg is in a tss but not in a gene
			if (arr[len - 4].compareTo("NA") != 0) {
				// CpG in at least one TSS
				for (int i = 0; i < tss_arr.length; i++) {
					// update tss status of the genes in gene list
					boolean is_gene_in_lst = false;
					for (GeneAnn gene : gene_lst) {
						if (gene.name.compareTo(tss_arr[i]) == 0) {
							// gene in gene list which overlaps this TSS
							is_gene_in_lst = true;
							gene.in_tss = true;
							break;
						}
					}

					// check if no gene found for this tss
					if (is_gene_in_lst == false) {
						// CpG in a TSS not overlapping a gene
						GeneAnn tss_gene = new GeneAnn(tss_arr[i], "NA");
						tss_gene.in_tss = true;
						gene_lst.add(tss_gene);
					}
				}
			}
		}
	}

	// gene class to handle tss, gene and exon/intron
	public static class GeneAnn implements Comparable<GeneAnn> {
		String name = "";
		boolean in_tss = false;
		String exon = "";

		// constructor
		public GeneAnn(String name, String exon) {
			this.name = name;
			this.exon = exon;
		}

		// compare the name of two genes
		public int compareTo(GeneAnn that) {
			return this.name.compareTo(that.name);
		}
	}

	// DMR gene class
	public static class GeneDMRAnn {
		String name;
		boolean in_tss = false;
		String exon;

		// constructor
		public GeneDMRAnn(GeneAnn gene) {
			this.name = gene.name;
			this.in_tss = gene.in_tss;
			this.exon = gene.exon;
		}

		// update annotation
		public void update(GeneAnn that) {
			// check
			if (this.name.compareTo(that.name) != 0) {
				System.err.println("gene names do not match " + this.name
						+ "\t" + that.name);
				return;
			}

			// update the tss status
			if (that.in_tss == true) {
				this.in_tss = true;
			}

			// update the exon status
			if (this.exon.compareTo("NA") == 0) {
				// not in gene body [but in tss], take any possible new value
				this.exon = that.exon;
			} else if (this.exon.compareTo("INTRON") == 0
					&& that.exon.compareTo("EXON") == 0) {
				// DMR in gene intron, update only if that is in an exon
				this.exon = that.exon;
			}
		}
	}
}
